use super::*;
use core::{ptr, slice, str};
use hierr::Error;

pub struct ElfDynFile<T: ElfBits> {
    fd: ElfFd,
    hash: T::SectionHdr,
    sym: T::SectionHdr,
    name: T::SectionHdr,
    len: usize,
}

#[allow(dead_code)]
impl<T: ElfBits> ElfDynFile<T> {
    pub fn open(fd: ElfFd) -> Result<Self, Error> {
        if !fd.is_bits::<T>() {
            return Err(Error::new(libc::EINVAL));
        }
        let len = fd.size();

        let mut hash = None;
        let mut sym = None;
        let mut name = None;
        for hdr in fd.section_iter::<T>() {
            match fd.read_section_name::<T>(hdr.sh_name()) {
                Some(".dynstr") => name = Some(hdr),
                Some(".gnu.hash") => hash = Some(hdr),
                Some(".dynsym") => sym = Some(hdr),
                _ => (),
            }
        }
        if hash.is_none() || sym.is_none() || name.is_none() {
            return Err(Error::new(libc::EINVAL));
        }
        Ok(Self {
            fd,
            hash: hash.unwrap(),
            sym: sym.unwrap(),
            name: name.unwrap(),
            len,
        })
    }

    pub fn as_bytes(&self) -> &[u8] {
        self.fd.read_slice::<u8>(0, self.len).unwrap()
    }

    pub fn sym_count(&self) -> u32 {
        (self.sym.sh_size() / self.sym.sh_entsize()) as u32
    }

    pub fn sym_name(&self, symndx: u32) -> Option<&str> {
        let sym = self.sym(symndx)?;
        if sym.st_name() as u64 >= self.name.sh_size() {
            return None;
        }
        let offset = self.name.sh_offset() as usize + sym.st_name() as usize;
        self.fd.read_cstr(offset)
    }

    fn sym(
        &self,
        symndx: u32,
    ) -> Option<T::Symbol> {
        if symndx >= self.sym_count() {
            return None;
        }
        let offset = self.sym.sh_offset() as usize
            + core::mem::size_of::<T::Symbol>()
                * symndx as usize;
        let sym = self
            .fd
            .read::<T::Symbol>(offset)?;
        Some(self.fd.map(sym))
    }

    fn sym_mut(
        &mut self,
        symndx: u32,
    ) -> Option<&mut T::Symbol> {
        if symndx >= self.sym_count() {
            return None;
        }
        let offset = self.sym.sh_offset() as usize
            + core::mem::size_of::<T::Symbol>()
                * symndx as usize;
        self.fd
            .read_mut::<T::Symbol>(offset)
    }

    fn hash_info(&self) -> Option<Elf_GnuHash> {
        let hdr = self
            .fd
            .read::<Elf_GnuHash>(self.hash.sh_offset() as usize)?;
        Some(self.fd.map(hdr))
    }

    pub fn size(&self) -> usize {
        self.len
    }

    pub fn hash_buckets(&self) -> u32 {
        self.hash_info().unwrap().gh_nbuckets
    }

    pub fn find(&self, name: &str) -> Option<u32> {
        let hdr = self.hash_info()?;
        let (hash, nbucket, bmndx, bmmask) = hdr.hash::<T>(name.as_bytes());
        let bmoff = hdr.bm_offset::<T>(bmndx) + self.hash.sh_offset() as usize;
        let bmfield = self.fd.map(self.fd.read::<T::Bits>(bmoff)?);
        if (bmfield & bmmask) != bmmask {
            return Some(u32::MAX);
        }
        let hboff = hdr.hb_offset::<T>(nbucket) + self.hash.sh_offset() as usize;
        let mut symndx = self.fd.map(self.fd.read::<u32>(hboff)?);

        let hash = hash >> 1;
        while symndx < self.sym_count() {
            let hvoff = hdr.hv_offset::<T>(symndx) + self.hash.sh_offset() as usize;
            let hval = self.fd.map(self.fd.read::<u32>(hvoff)?);
            if (hval >> 1) == hash && self.sym_name(symndx)? == name {
                return Some(symndx);
            }
            if (hval & 0x01) == 0x01 {
                break;
            }
            symndx += 1;
        }
        Some(u16::MAX as u32)
    }

    pub fn export_iter(&self) -> impl Iterator<Item = (u32, &str)> {
        DynSymIter::new(
            self,
            self.hash_info().map_or(self.sym_count(), |h| h.gh_symndx),
            self.sym_count(),
        )
    }

    pub fn import_iter(&self) -> impl Iterator<Item = (u32, &str)> {
        DynSymIter::new(self, 1, self.hash_info().map_or(1, |h| h.gh_symndx))
    }

    pub fn new_name_buff(&self) -> Result<NameBuff, Error> {
        NameBuff::new::<T>(self)
    }

    pub fn update_merge(&mut self, buff: &mut NameBuff) -> u32 {
        let diff = self.update_name(buff);
        if diff == 0 {
            return 0;
        }
        self.update_hash(buff);
        self.update_offset(diff);
        self.update_others(diff);

        diff
    }

    fn update_others(&mut self, diff: u32) {
        self.len -= diff as usize;

        let off1 = (self.name.sh_offset() + self.name.sh_size()) as usize;
        let off2 = off1 + diff as usize;
        let src = self.read_slice::<u8>(off2, 1).unwrap().as_ptr();
        let dst = self.read_slice_mut::<u8>(off1, 1).unwrap().as_mut_ptr();
        unsafe { ptr::copy(src, dst, self.len - off1) };
    }

    fn update_hash(&mut self, buff: &mut NameBuff) {
        let info = self.hash_info().unwrap();
        let host_endian = self.host_endian();

        let bmoff = info.bm_offset::<T>(0) + self.hash.sh_offset() as usize;
        let masks = self
            .read_slice_mut::<T::Bits>(bmoff, info.gh_bmmask as usize)
            .unwrap();
        for m in masks {
            *m = <T::Bits>::default();
        }

        for n in info.gh_symndx..self.sym_count() {
            let name = buff.name_by_symndx(n).unwrap();
            let (hash, _, bmndx, bmmask) = info.hash::<T>(name.as_bytes());
            //new hash value
            let hash = hash.map(self.host_endian());
            let voff = info.hv_offset::<T>(n) + self.hash.sh_offset() as usize;
            let old = self.read::<u32>(voff).unwrap();
            let hash = (hash & !0x01) | (old & 0x01);
            *self.fd.read_mut::<u32>(voff).unwrap() = hash.map(host_endian);
            //new bmmask value
            let bmoff = info.bm_offset::<T>(bmndx) + self.hash.sh_offset() as usize;
            let bmval = self.read::<T::Bits>(bmoff).unwrap().map(host_endian) | bmmask;
            *self.read_mut::<T::Bits>(bmoff).unwrap() = bmval.map(host_endian);
        }
    }

    fn update_name(&mut self, buff: &mut NameBuff) -> u32 {
        let align = self.max_align();
        let new = buff.merge(align);

        let offset = self.name.sh_offset() as usize;
        let old = self.read_slice_mut::<u8>(offset, new.len()).unwrap();
        old.copy_from_slice(new);

        let diff = self.name.sh_size() as u32 - new.len() as u32;
        self.name.sh_size_set(new.len() as u64);

        let host_endian = self.host_endian();
        let name_offset = self.name.sh_offset().map(host_endian);
        let name_size = self.name.sh_size().map(host_endian);
        for hdr in self.section_iter_mut::<T>() {
            if hdr.sh_offset() == name_offset {
                hdr.sh_size_set(name_size);
                break;
            }
        }

        for n in 1..self.sym_count() {
            let sym = self.sym_mut(n).unwrap();
            sym.st_name_set(buff.name_offset(n).unwrap().map(host_endian));
        }

        diff
    }

    fn max_align(&mut self) -> u32 {
        let mut max_align = 1;
        for hdr in self.section_iter::<T>() {
            if hdr.sh_offset() > self.name.sh_offset() && hdr.sh_addralign() > max_align {
                max_align = hdr.sh_addralign();
            }
        }
        max_align as u32
    }

    fn set<F>(min: u64, host_endian: bool, diff: u32, new: u64, mut f: F)
    where
        F: FnMut(u64),
    {
        let new = new.map(host_endian);
        if new > min {
            f((new - diff as u64).map(host_endian));
        }
    }

    fn update_offset(&mut self, diff: u32) {
        let name_offset = self.name.sh_offset();
        let name_addr = self.name.sh_addr() + self.name.sh_size();
        let mut end_addr = u64::MAX;
        let host_endian = self.host_endian();

        // 首先查找到对应的Segment, 修改其大小
        // 只有Segment内部大于name_addr的才会被修改
        for hdr in self.segment_iter_mut::<T>() {
            let offset = hdr.ph_offset().map(host_endian);
            let filesz = hdr.ph_filesz().map(host_endian);
            if offset <= name_offset && offset + filesz > name_offset {
                hdr.ph_filesz_set((filesz - diff as u64).map(host_endian));
                let memsz = hdr.ph_memsz().map(host_endian);
                hdr.ph_memsz_set((memsz - diff as u64).map(host_endian));
                //是否需要限定修改地址的范围???
                //let vaddr = hdr.ph_vaddr().map(host_endian);
                //end_addr = vaddr + memsz;
                break;
            }
        }

        for hdr in self.segment_iter_mut::<T>() {
            Self::set(name_offset, host_endian, diff, hdr.ph_offset(), |off| {
                hdr.ph_offset_set(off)
            });
                Self::set(name_addr, host_endian, diff, hdr.ph_vaddr(), |addr| {
                    hdr.ph_vaddr_set(addr)
                });
                Self::set(name_addr, host_endian, diff, hdr.ph_paddr(), |addr| {
                    hdr.ph_paddr_set(addr)
                });
        }

        // 首先调整.dynamic中的地址
        for e in self.dyn_iter_mut::<T>() {
                    let addr = e.d_val_ptr().map(host_endian);
                    let tag = e.d_tag().map(host_endian);
                    if d_tag_is_ptr(tag) && addr > name_addr && addr < end_addr {
                        let addr = addr - diff as u64;
                        e.d_val_ptr_set(addr.map(host_endian));
                    }
        }

        // 最后调整所有section的偏移和addr
        for hdr in self.section_iter_mut::<T>() {
            let addr = hdr.sh_addr().map(host_endian);
            if addr < end_addr {
                Self::set(name_addr, host_endian, diff, hdr.sh_addr(), |addr| {
                    hdr.sh_addr_set(addr)
                });
            }
            Self::set(name_offset, host_endian, diff, hdr.sh_offset(), |offset| {
                hdr.sh_offset_set(offset)
            });
        }

        let ehdr = self.read_mut::<T::ElfHdr>(0).unwrap();
        Self::set(name_offset, host_endian, diff, ehdr.e_shoff(), |off| {
            ehdr.e_shoff_set(off)
        });
        Self::set(name_offset, host_endian, diff, ehdr.e_phoff(), |off| {
            ehdr.e_phoff_set(off)
        });

        Self::set(name_offset, true, diff, self.sym.sh_offset(), |off| {
            self.sym.sh_offset_set(off)
        });
        Self::set(name_offset, true, diff, self.hash.sh_offset(), |off| {
            self.hash.sh_offset_set(off)
        });
    }
}

struct DynSymIter<'a, T: ElfBits> {
    fd: &'a ElfDynFile<T>,
    ndx: u32,
    end: u32,
}

impl<'a, T: ElfBits> DynSymIter<'a, T> {
    fn new(fd: &'a ElfDynFile<T>, ndx: u32, end: u32) -> Self {
        Self { fd, ndx, end }
    }
}

impl<'a, T: ElfBits> Iterator for DynSymIter<'a, T> {
    type Item = (u32, &'a str);
    fn next(&mut self) -> Option<Self::Item> {
        if self.ndx < self.end {
            let ndx = self.ndx;
            self.ndx += 1;
            Some((ndx, self.fd.sym_name(ndx).unwrap()))
        } else {
            None
        }
    }
}

impl<T: ElfBits> core::ops::Deref for ElfDynFile<T> {
    type Target = ElfFd;
    fn deref(&self) -> &Self::Target {
        &self.fd
    }
}

impl<T: ElfBits> core::ops::DerefMut for ElfDynFile<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.fd
    }
}

pub struct NameBuff {
    nndx: *mut u32,
    symcount: u32,
    name: *mut u8,
    namesize: u32,
}

impl NameBuff {
    fn new<T: ElfBits>(file: &ElfDynFile<T>) -> Result<Self, Error> {
        let symcount = file.sym_count();
        let namesize = file.name.sh_size() as u32;
        let capacity = symcount as usize * core::mem::size_of::<u32>() + namesize as usize;
        let nndx = unsafe { libc::malloc(capacity) };
        if nndx.is_null() {
            return Err(Error::new(libc::ENOMEM));
        }
        unsafe { libc::memset(nndx, 0, capacity) };
        let nndx = nndx.cast::<u32>();

        for n in 0..symcount {
            let sym = file.sym(n).unwrap();
            unsafe { nndx.add(n as usize).write(sym.st_name()) };
        }

        let name = unsafe { nndx.add(symcount as usize).cast::<u8>() };
        let name_src = file
            .read_slice::<u8>(file.name.sh_offset() as usize, namesize as usize)
            .ok_or(Error::default())?;
        unsafe { ptr::copy_nonoverlapping(name_src.as_ptr(), name, namesize as usize) };

        Ok(Self {
            nndx,
            symcount,
            name,
            namesize,
        })
    }

    // 如果符号名存在重合，则不支持修改
    pub fn write(&mut self, nbuckets: u32, symndx: u32, new: &str) -> bool {
        let suffix = &mut [0_u8; 8];
        let Some(offset) = self.name_offset(symndx) else {
            return false;
        };
        let Some(old) = self.name_by_offset(offset) else {
            return false;
        };
        if !self.pre_0(offset) {
            return false;
        }
        if self.name_offset_in(offset, offset + old.len() as u32) {
            return false;
        }

        let Some(suffix) = new_name_suffix(new, old, nbuckets, suffix) else {
            return false;
        };

        let suffix_0 = self.suffix_0_count(offset + old.len() as u32) as usize;
        if new.len() + suffix.len() >= old.len() + suffix_0 - 1 {
            return false;
        }

        let dst = unsafe { self.name.add(offset as usize) };
        unsafe { ptr::copy_nonoverlapping(new.as_ptr(), dst, new.len()) };
        let dst = unsafe { dst.add(new.len()) };
        unsafe { ptr::copy_nonoverlapping(suffix.as_ptr(), dst, suffix.len()) };
        let dst = unsafe { dst.add(suffix.len()) };

        for n in 0..old.len() - new.len() - suffix.len() {
            unsafe { dst.add(n).write(0) };
        }

        true
    }

    pub fn restore(&mut self, symndx: u32, new: &str) -> bool {
        let Some(offset) = self.name_offset(symndx) else {
            return false;
        };
        let Some(old) = self.name_by_offset(offset) else {
            return false;
        };
        if new.len() < old.len() {
            return false;
        }

        let suffix_0 = self.suffix_0_count(offset + old.len() as u32) as usize;
        if new.len() > old.len() + suffix_0 - 1 {
            return false;
        }

        let dst = unsafe { self.name.add(offset as usize) };
        unsafe { ptr::copy_nonoverlapping(new.as_ptr(), dst, new.len()) };

        true
    }

    fn name_offset(&self, symndx: u32) -> Option<u32> {
        if symndx >= self.symcount {
            return None;
        }
        Some(unsafe { self.nndx.add(symndx as usize).read() })
    }

    pub fn name_by_symndx(&self, symndx: u32) -> Option<&str> {
        let offset = self.name_offset(symndx)?;
        self.name_by_offset(offset)
    }

    fn name_by_offset(&self, offset: u32) -> Option<&str> {
        for n in offset..self.namesize {
            unsafe {
                if self.name.add(n as usize).read() == 0 {
                    let buf = slice::from_raw_parts(
                        self.name.add(offset as usize),
                        (n - offset) as usize,
                    );
                    return Some(str::from_utf8_unchecked(buf));
                }
            }
        }
        None
    }

    // 考虑到可能存在符号名重合的场景
    fn merge(&mut self, max_align: u32) -> &[u8] {
        let mut diff = self.last_0_count() - 1;
        for n in 1..self.symcount {
            let offset = self.name_offset(n).unwrap();
            let pre_0 = self.pre_0_count(offset);
            if pre_0 <= 1 {
                continue;
            }
            self.move_forward(offset, pre_0 - 1);
            diff += pre_0 - 1;
        }

        let aligned_diff = diff & !(max_align - 1);
        for n in aligned_diff..diff {
            unsafe { self.name.add(n as usize).write(0) };
        }

        let len = self.namesize - aligned_diff;
        unsafe { core::slice::from_raw_parts(self.name, len as usize) }
    }

    fn move_forward(&mut self, offset: u32, distance: u32) {
        let src = unsafe { self.name.add(offset as usize) };
        let dst = unsafe { src.sub(distance as usize) };
        let len = self.namesize - offset;
        unsafe { ptr::copy(src, dst, len as usize) };

        for n in 1..self.symcount {
            let nndx = unsafe { self.nndx.add(n as usize) };
            let val = unsafe { nndx.read() };
            if val >= offset {
                unsafe { nndx.write(val - distance) };
            }
        }
    }

    fn pre_0_count(&self, offset: u32) -> u32 {
        let mut cnt = 0;
        for n in (0..offset).rev() {
            let c = unsafe { self.name.add(n as usize).read() };
            if c != 0 {
                break;
            }
            cnt += 1;
        }
        cnt
    }

    fn suffix_0_count(&self, offset: u32) -> u32 {
        let mut cnt = 0;
        for n in offset..self.namesize {
            let c = unsafe { self.name.add(n as usize).read() };
            if c != 0 {
                break;
            }
            cnt += 1;
        }
        cnt
    }

    fn last_0_count(&self) -> u32 {
        self.pre_0_count(self.namesize)
    }

    fn pre_0(&self, offset: u32) -> bool {
        unsafe { self.name.add(offset as usize - 1).read() == 0 }
    }

    fn name_offset_in(&self, beg: u32, end: u32) -> bool {
        for n in 1..self.symcount {
            let offset = unsafe { self.nndx.add(n as usize).read() };
            if offset > beg && offset < end {
                return true;
            }
        }
        false
    }
}

impl Drop for NameBuff {
    fn drop(&mut self) {
        unsafe { libc::free(self.nndx.cast::<libc::c_void>()) };
    }
}

fn new_name_suffix<'a>(
    new: &str,
    old: &str,
    nbuckets: u32,
    suffix: &'a mut [u8],
) -> Option<&'a str> {
    let old = gnu_hash(old.as_bytes()) % nbuckets;
    let new_hash = gnu_hash(new.as_bytes());
    let chars = b"$._0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz";
    let mut len = 0;
    if (new_hash % nbuckets) == old {
        return Some("");
    }
    'top: while len < suffix.len() {
        let mut new = new_hash;
        for n in 0..len {
            let c = chars[suffix[n] as usize] as u32;
            new = new.wrapping_mul(33) + c;
        }
        for c in chars {
            let hash = new.wrapping_mul(33) + *c as u32;
            if (hash % nbuckets) == old {
                for n in 0..len {
                    suffix[n] = chars[suffix[n] as usize];
                }
                suffix[len] = *c;
                return Some(unsafe { core::str::from_utf8_unchecked(&suffix[..len + 1]) });
            }
        }
        for n in (0..len).rev() {
            suffix[n] += 1;
            if (suffix[n] as usize) < chars.len() {
                continue 'top;
            }
            suffix[n] = 0;
        }
        suffix[len] = 0;
        len += 1;
    }
    None
}
